from os.path import splitext
from os import listdir
import numpy as np
import cv2
from glob import glob
import torch
from torch.utils.data import Dataset
import logging
from PIL import Image
import os


class BasicDataset(Dataset):
    def __init__(self, imgs_dir, masks_dir, scale=1, mask_suffix=''):
        self.images_path = []
        self.labels_path = []
        files = os.listdir(imgs_dir)
        for file in files:
            basename = os.path.basename(file).split('.')[0]
            self.images_path.append(imgs_dir + basename + '.jpg')
            self.labels_path.append(masks_dir + basename + '.png')
        self.len = len(self.images_path)
        self.size = (512, 512)
        logging.info(f'Creating dataset with {self.len} examples')

    def __len__(self):
        return self.len

    @classmethod
    def preprocess(cls, pil_img, size):
        w, h = size
        assert w > 100 and h > 100, 'Image is too small'
        img_nd = np.array(pil_img)
        img_nd = cv2.resize(img_nd, size, interpolation=cv2.INTER_NEAREST)

        if len(img_nd.shape) == 2:
            img_nd = np.expand_dims(img_nd, axis=2)

        # HWC to CHW
        img_trans = img_nd.transpose((2, 0, 1))

        return img_trans

    def __getitem__(self, i):
        img = Image.open(self.images_path[i])
        mask = Image.open(self.labels_path[i])

        img = self.preprocess(img, self.size)
        mask = self.preprocess(mask, self.size)

        classes = np.unique(mask).astype(np.int)
        assert len(classes) > 2

        return {
            'image': torch.from_numpy(img).float(),
            'mask': torch.from_numpy(mask).float()
        }


class CarvanaDataset(BasicDataset):
    def __init__(self, imgs_dir, masks_dir, scale=1):
        super().__init__(imgs_dir, masks_dir, scale, mask_suffix='_mask')
